package com.iakuil.app.common.ws.utils;

import com.iakuil.app.common.core.constant.RedisConstant;
import com.iakuil.app.common.core.util.SpinLockUtils;
import lombok.experimental.UtilityClass;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * Websocket Session工具类
 *
 * <p>对本地Websocket Session列表进行操作。
 *
 * @author Kai
 */
@Slf4j
@UtilityClass
public class WsSessionUtils {
    /**
     * 当前节点在线session
     */
    public static Map<String, WebSocketSession> CURRENT_NODE_SESSIONS = new ConcurrentHashMap<>();

    /**
     * WS客户端接入
     *
     * @param session WS会话对象
     * @return 当前节点客户端连接数
     */
    public static int join(WebSocketSession session) {
        String uid = String.valueOf(session.getAttributes().get(RedisConstant.WS_ID_ATTR_NAME));
        CURRENT_NODE_SESSIONS.put(uid, session);
        log.info("建立WS连接：[{}]，来自[{}]", uid, session.getUri());
        return CURRENT_NODE_SESSIONS.size();
    }

    /**
     * WS客户端离线
     *
     * @param session WS会话对象
     * @return 当前节点客户端连接数
     */
    public static int leave(WebSocketSession session) {
        String uid = String.valueOf(session.getAttributes().get(RedisConstant.WS_ID_ATTR_NAME));
        CURRENT_NODE_SESSIONS.remove(uid);
        log.info("断开WS连接: uid[{}]", uid);
        return CURRENT_NODE_SESSIONS.size();
    }

    /**
     * 当前节点是否包含会话ID
     *
     * @param sessionId WS会话ID
     * @return 是否当前节点会话
     */
    public static boolean contains(String sessionId) {
        return CURRENT_NODE_SESSIONS.containsKey(sessionId);
    }

    /**
     * 群发给当前节点所有客户端
     *
     * @param msg 文本消息
     */
    public static void sendToAll(String msg) {
        for (WebSocketSession session : CURRENT_NODE_SESSIONS.values()) {
            if (session.isOpen()) {
                doSend(session, msg);
            }
        }
    }

    /**
     * 发送消息给指定客户端
     *
     * <p>不是本地节点会话直接忽略。
     *
     * @param toId 接收人会话ID
     * @param msg  文本消息
     */
    public static void sendMessage(String toId, String msg) {
        WebSocketSession session = CURRENT_NODE_SESSIONS.get(toId);
        if (session != null && session.isOpen()) {
            doSend(session, msg);
        }
    }

    private static void doSend(WebSocketSession session, String msg) {
        try {
            SpinLockUtils.lock(session, false);
            session.sendMessage(new TextMessage(msg));
        } catch (IOException e) {
            throw new IllegalStateException("Occurring an exception during message sending through websocket!", e);
        } finally {
            SpinLockUtils.unlock(session);
        }
    }
}
